package resml

import (
	"io"
	"sync"
	"time"
)

// 具备写入进度能力的 io.Writer
type ProgressWriter struct {
	mu    sync.Mutex
	w     io.Writer // 底层的io.Writer
	n     int64     // 已写入的字节数
	total int64     // 总的字节数
}

// 创建一个 ProgressWriter
func NewProgressWriter(w io.Writer, total int64) *ProgressWriter {
	return &ProgressWriter{w: w, total: total}
}

// 实现 io.Writer
func (w *ProgressWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	n, err := w.w.Write(p)
	w.n += int64(n)
	return n, err
}

// 进度百分比，当总量值为 0 时，返回 0
func (w *ProgressWriter) Percent() float64 {
	if w.total == 0 {
		return 0
	}
	return float64(w.n) / float64(w.total) * 100
}

// 当前进度
func (w *ProgressWriter) Current() int64 {
	return w.n
}

// 总量值
func (w *ProgressWriter) Total() int64 {
	return w.total
}

// 限定写入字节数的 io.Writer
type LimitedWriter struct {
	mu sync.Mutex
	w  io.Writer // 底层的io.Writer
	n  int64     // 已写入的字节数
	l  int64     // 限制写入的字节数
}

// 创建一个 LimitedWriter
func NewLimitedWriter(w io.Writer, limit int64) *LimitedWriter {
	return &LimitedWriter{w: w, l: limit}
}

// 实现 io.Writer
func (w *LimitedWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.n >= w.l {
		return len(p), nil // 已写满，丢弃并返回
	}
	s := int64(len(p))
	if w.n+s > w.l {
		tbw := w.l - w.n     // to be writed
		dis := w.n + s - w.l // discard
		n, err := w.w.Write(p[:tbw])
		w.n += int64(n)
		return n + int(dis), err // 丢弃掉的被视为已经写入成功的
	}
	n, err := w.w.Write(p)
	w.n += int64(n)
	return n, err
}

// RateLimitedWriter 限速写入器。可用于测试场景。
type RateLimitedWriter struct {
	mu    sync.Mutex
	w     io.Writer // 底层的 writer
	start time.Time // 开始时间
	rate  int64     // 每秒写入的字节数
	n     int64     // 已经写入的字节数
}

func NewRateLimitedWriter(w io.Writer, bytesPerSecond int64) *RateLimitedWriter {
	return &RateLimitedWriter{w: w, rate: bytesPerSecond}
}

// 实现 io.Writer
func (w *RateLimitedWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.start.IsZero() {
		w.start = time.Now()
	}
	td := w.since()
	for (w.n / int64(td.Seconds())) > w.rate {
		time.Sleep(time.Second)
		td = w.since()
	}
	n, err := w.w.Write(p)
	w.n += int64(n)
	return n, err
}

func (w *RateLimitedWriter) since() time.Duration {
	d := time.Since(w.start)
	if d < time.Second {
		return time.Second
	}
	return d
}
